import pandas as pd
from snownlp import SnowNLP

# 定义一个函数来判断新闻类型
def classify_news(news_content):
    # 使用 SnowNLP 判断情感倾向
    s = SnowNLP(news_content)
    sentiment = s.sentiments

    # 判断与中国或美国的相关性
    if '中国' in news_content:
        if sentiment >= 0.6:  # 情感积极
            return '对中国正面'
        elif sentiment <= 0.4:  # 情感消极
            return '对中国负面'
        else:
            return '与中美无关'
    elif '美国' in news_content:
        if sentiment >= 0.6:
            return '对美国正面'
        elif sentiment <= 0.4:
            return '对美国负面'
        else:
            return '与中美无关'
    else:
        return '与中美无关'

# 加载数据
df = pd.read_csv('./chinanews_news.csv')

# 对每一条新闻内容进行类型判断，并插入到第二列
df.insert(1, '新闻类型', df['news_content'].apply(classify_news))

# 将结果保存为新的 CSV 文件
new_csv_path = './chinanews_news_classified_snowNLP.csv'
df.to_csv(new_csv_path, index=False)